import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, repeat  # 用于维度重排的工具库

import math

class DSW_embedding(nn.Module):
    def __init__(self, seg_len, d_model):
        super(DSW_embedding, self).__init__()
        self.seg_len = seg_len    # 每个分段的长度（如 6 个时间步为一段）

        self.linear = nn.Linear(seg_len, d_model)    # 线性层：将分段映射到 d_model 维度

    def forward(self, x):
        # 输入 x 形状：[batch_size, 时间步长度, 数据维度]（例如 [32, 168, 7]）
        batch, ts_len, ts_dim = x.shape

        # 步骤1：将时间序列按 seg_len 分段
        # 重排后形状：[(batch * 数据维度 * 分段数), seg_len]
        x_segment = rearrange(x, 'b (seg_num seg_len) d -> (b d seg_num) seg_len', seg_len = self.seg_len)
        # 步骤2：通过线性层将每个分段映射为 d_model 维度的特征
        x_embed = self.linear(x_segment)
        # 步骤3：恢复维度顺序，便于后续处理
        # 输出形状：[batch_size, 数据维度, 分段数, d_model]（例如 [32, 7, 28, 256]）
        x_embed = rearrange(x_embed, '(b d seg_num) d_model -> b d seg_num d_model', b = batch, d = ts_dim)
        
        return x_embed